home *** CD-ROM | disk | FTP | other *** search
/ Personal Computer World 2008 February / PCWFEB08.iso / Software / Freeware / Miro 1.0 / Miro_Installer.exe / xulrunner / python / BitTorrent / Uploader.py < prev    next >
Encoding:
Python Source  |  2007-11-12  |  7.9 KB  |  278 lines
  1. # Written by Bram Cohen
  2. # see LICENSE.txt for license information
  3.  
  4. from CurrentRateMeasure import Measure
  5.  
  6. class Upload:
  7.     def __init__(self, connection, choker, storage, 
  8.             max_slice_length, max_rate_period, fudge):
  9.         self.connection = connection
  10.         self.choker = choker
  11.         self.storage = storage
  12.         self.max_slice_length = max_slice_length
  13.         self.max_rate_period = max_rate_period
  14.         self.choked = True
  15.         self.interested = False
  16.         self.buffer = []
  17.         self.measure = Measure(max_rate_period, fudge)
  18.         if storage.do_I_have_anything():
  19.             connection.send_bitfield(storage.get_have_list())
  20.  
  21.     def got_not_interested(self):
  22.         if self.interested:
  23.             self.interested = False
  24.             del self.buffer[:]
  25.             self.choker.not_interested(self.connection)
  26.  
  27.     def got_interested(self):
  28.         if not self.interested:
  29.             self.interested = True
  30.             self.choker.interested(self.connection)
  31.  
  32.     def flushed(self):
  33.         while len(self.buffer) > 0 and self.connection.is_flushed():
  34.             index, begin, length = self.buffer[0]
  35.             del self.buffer[0]
  36.             piece = self.storage.get_piece(index, begin, length)
  37.             if piece is None:
  38.                 self.connection.close()
  39.                 return
  40.             self.measure.update_rate(len(piece))
  41.             self.connection.send_piece(index, begin, piece)
  42.  
  43.     def got_request(self, index, begin, length):
  44.         if not self.interested or length > self.max_slice_length:
  45.             self.connection.close()
  46.             return
  47.         if not self.choked:
  48.             self.buffer.append((index, begin, length))
  49.             self.flushed()
  50.  
  51.     def got_cancel(self, index, begin, length):
  52.         try:
  53.             self.buffer.remove((index, begin, length))
  54.         except ValueError:
  55.             pass
  56.  
  57.     def choke(self):
  58.         if not self.choked:
  59.             self.choked = True
  60.             del self.buffer[:]
  61.             self.connection.send_choke()
  62.  
  63.     def unchoke(self):
  64.         if self.choked:
  65.             self.choked = False
  66.             self.connection.send_unchoke()
  67.         
  68.     def is_choked(self):
  69.         return self.choked
  70.         
  71.     def is_interested(self):
  72.         return self.interested
  73.  
  74.     def has_queries(self):
  75.         return len(self.buffer) > 0
  76.  
  77.     def get_rate(self):
  78.         return self.measure.get_rate()
  79.  
  80. class DummyConnection:
  81.     def __init__(self, events):
  82.         self.events = events
  83.         self.flushed = False
  84.  
  85.     def send_bitfield(self, bitfield):
  86.         self.events.append(('bitfield', bitfield))
  87.     
  88.     def is_flushed(self):
  89.         return self.flushed
  90.  
  91.     def close(self):
  92.         self.events.append('closed')
  93.  
  94.     def send_piece(self, index, begin, piece):
  95.         self.events.append(('piece', index, begin, piece))
  96.  
  97.     def send_choke(self):
  98.         self.events.append('choke')
  99.  
  100.     def send_unchoke(self):
  101.         self.events.append('unchoke')
  102.  
  103. class DummyChoker:
  104.     def __init__(self, events):
  105.         self.events = events
  106.  
  107.     def interested(self, connection):
  108.         self.events.append('interested')
  109.     
  110.     def not_interested(self, connection):
  111.         self.events.append('not interested')
  112.  
  113. class DummyStorage:
  114.     def __init__(self, events):
  115.         self.events = events
  116.  
  117.     def do_I_have_anything(self):
  118.         self.events.append('do I have')
  119.         return True
  120.  
  121.     def get_have_list(self):
  122.         self.events.append('get have list')
  123.         return [False, True]
  124.  
  125.     def get_piece(self, index, begin, length):
  126.         self.events.append(('get piece', index, begin, length))
  127.         if length == 4:
  128.             return None
  129.         return 'a' * length
  130.  
  131. def test_skip_over_choke():
  132.     events = []
  133.     dco = DummyConnection(events)
  134.     dch = DummyChoker(events)
  135.     ds = DummyStorage(events)
  136.     u = Upload(dco, dch, ds, 100, 20, 5)
  137.     assert u.is_choked()
  138.     assert not u.is_interested()
  139.     u.got_interested()
  140.     assert u.is_interested()
  141.     u.got_request(0, 0, 3)
  142.     dco.flushed = True
  143.     u.flushed()
  144.     assert events == ['do I have', 'get have list', 
  145.         ('bitfield', [False, True]), 'interested']
  146.  
  147. def test_bad_piece():
  148.     events = []
  149.     dco = DummyConnection(events)
  150.     dch = DummyChoker(events)
  151.     ds = DummyStorage(events)
  152.     u = Upload(dco, dch, ds, 100, 20, 5)
  153.     assert u.is_choked()
  154.     assert not u.is_interested()
  155.     u.got_interested()
  156.     assert u.is_interested()
  157.     u.unchoke()
  158.     assert not u.is_choked()
  159.     u.got_request(0, 0, 4)
  160.     dco.flushed = True
  161.     u.flushed()
  162.     assert events == ['do I have', 'get have list', 
  163.         ('bitfield', [False, True]), 'interested', 'unchoke', 
  164.         ('get piece', 0, 0, 4), 'closed']
  165.  
  166. def test_still_rejected_after_unchoke():
  167.     events = []
  168.     dco = DummyConnection(events)
  169.     dch = DummyChoker(events)
  170.     ds = DummyStorage(events)
  171.     u = Upload(dco, dch, ds, 100, 20, 5)
  172.     assert u.is_choked()
  173.     assert not u.is_interested()
  174.     u.got_interested()
  175.     assert u.is_interested()
  176.     u.unchoke()
  177.     assert not u.is_choked()
  178.     u.got_request(0, 0, 3)
  179.     u.choke()
  180.     u.unchoke()
  181.     dco.flushed = True
  182.     u.flushed()
  183.     assert events == ['do I have', 'get have list', 
  184.         ('bitfield', [False, True]), 'interested', 'unchoke', 
  185.         'choke', 'unchoke']
  186.  
  187. def test_sends_when_flushed():
  188.     events = []
  189.     dco = DummyConnection(events)
  190.     dch = DummyChoker(events)
  191.     ds = DummyStorage(events)
  192.     u = Upload(dco, dch, ds, 100, 20, 5)
  193.     u.unchoke()
  194.     u.got_interested()
  195.     u.got_request(0, 1, 3)
  196.     dco.flushed = True
  197.     u.flushed()
  198.     u.flushed()
  199.     assert events == ['do I have', 'get have list', 
  200.         ('bitfield', [False, True]), 'unchoke', 'interested', 
  201.         ('get piece', 0, 1, 3), ('piece', 0, 1, 'aaa')]
  202.  
  203. def test_sends_immediately():
  204.     events = []
  205.     dco = DummyConnection(events)
  206.     dch = DummyChoker(events)
  207.     ds = DummyStorage(events)
  208.     u = Upload(dco, dch, ds, 100, 20, 5)
  209.     u.unchoke()
  210.     u.got_interested()
  211.     dco.flushed = True
  212.     u.got_request(0, 1, 3)
  213.     assert events == ['do I have', 'get have list', 
  214.         ('bitfield', [False, True]), 'unchoke', 'interested', 
  215.         ('get piece', 0, 1, 3), ('piece', 0, 1, 'aaa')]
  216.  
  217. def test_cancel():
  218.     events = []
  219.     dco = DummyConnection(events)
  220.     dch = DummyChoker(events)
  221.     ds = DummyStorage(events)
  222.     u = Upload(dco, dch, ds, 100, 20, 5)
  223.     u.unchoke()
  224.     u.got_interested()
  225.     u.got_request(0, 1, 3)
  226.     u.got_cancel(0, 1, 3)
  227.     u.got_cancel(0, 1, 2)
  228.     u.flushed()
  229.     dco.flushed = True
  230.     assert events == ['do I have', 'get have list', 
  231.         ('bitfield', [False, True]), 'unchoke', 'interested']
  232.  
  233. def test_clears_on_not_interested():
  234.     events = []
  235.     dco = DummyConnection(events)
  236.     dch = DummyChoker(events)
  237.     ds = DummyStorage(events)
  238.     u = Upload(dco, dch, ds, 100, 20, 5)
  239.     u.unchoke()
  240.     u.got_interested()
  241.     u.got_request(0, 1, 3)
  242.     u.got_not_interested()
  243.     dco.flushed = True
  244.     u.flushed()
  245.     assert events == ['do I have', 'get have list', 
  246.         ('bitfield', [False, True]), 'unchoke', 'interested', 
  247.         'not interested']
  248.  
  249. def test_close_when_sends_on_not_interested():
  250.     events = []
  251.     dco = DummyConnection(events)
  252.     dch = DummyChoker(events)
  253.     ds = DummyStorage(events)
  254.     u = Upload(dco, dch, ds, 100, 20, 5)
  255.     u.got_request(0, 1, 3)
  256.     assert events == ['do I have', 'get have list', 
  257.         ('bitfield', [False, True]), 'closed']
  258.  
  259. def test_close_over_max_length():
  260.     events = []
  261.     dco = DummyConnection(events)
  262.     dch = DummyChoker(events)
  263.     ds = DummyStorage(events)
  264.     u = Upload(dco, dch, ds, 100, 20, 5)
  265.     u.got_interested()
  266.     u.got_request(0, 1, 101)
  267.     assert events == ['do I have', 'get have list', 
  268.         ('bitfield', [False, True]), 'interested', 'closed']
  269.  
  270. def test_no_bitfield_on_start_empty():
  271.     events = []
  272.     dco = DummyConnection(events)
  273.     dch = DummyChoker(events)
  274.     ds = DummyStorage(events)
  275.     ds.do_I_have_anything = lambda: False
  276.     u = Upload(dco, dch, ds, 100, 20, 5)
  277.     assert events == []
  278.